'''
Created on Mar 3, 2011

@author: oabalbin
'''
import os
import numpy as np
import scipy as sp
from scipy import stats
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import matplotlib.pylab as pla
import matplotlib.cm as cm 


def plot_scatter(mat1,mat2, corr, pli, folderpath,geneName=[]):
    """
    outa[0,0],outa[0,1],outa[0,3],outa[0,4], outa[0,5] = i (mat1 row index), k(factor number), pval, pcor, bpval, int(cmpname[i]) 
    (drug index in the NCI database).
    """    
    '''Linear regression'''
    
    #mat2 = 2*(mat2/float((np.max(mat2) - np.min(mat2))))
    xmin,xmax,ymin,ymax = np.min(mat1),np.max(mat1),np.min(mat2),np.max(mat2)
    
    # Linear regression    
    x = np.linspace(xmin,xmax,mat1.size)
    (ar,br) = np.polyfit(x,mat2,1)
    yn=np.polyval([ar,br],x)
    
    mat1sd = np.std(mat1)
    mat2sd = np.std(mat2)

    plt.xlim( (xmin-mat1sd, xmax+mat1sd) )
    plt.ylim( (ymin-0.5, ymax+0.5) )
    #plt.title('Linear regression with polifit()')
    #plt.plot(x,yn,'b-') #label='linear regression'
    plt.plot(mat1,mat2, 'ko')
    #legend(loc='best')
        
    plt.ylabel('Pathway activity (score)')
    plt.xlabel('Drug sensitivity (GI50) pval=%.4f' % round(corr[4],4) + ' R=%.2f' % round(corr[3],2))
    #xlim( (xmin, xmax) )     
    #plt.show()
    # Remember that not intercept is in the factors so the actual number of the factors with BFRM
    # Corresponds to j+1. corr[5] corresponds to the index of the drug in the NCI drug database.
    
    
    if geneName:
        pltname=os.path.join(folderpath,'plt_'+str(corr[5])+'_'+str(geneName[int(corr[1])])+'_'+str(pli)+'_.pdf')
    else:
        pltname=os.path.join(folderpath,'plt_'+str(corr[5])+'_'+str(corr[1] + 1)+'.pdf')

    plt.savefig(pltname)
    # Change to use with the general progam
    #plt.savefig(folderpath+'plt_'+str(corr[5])+'_'+str(corr[1])+'_'+'.png')
    plt.close()
    


def plot_scatter2(mat1,mat2, figname, folderpath, labels={}):
    
    plt.scatter(mat1, mat2, c='k',s=2, marker='o')
    
    #plt.savefig(folderpath+'plt_'+str(corr[0])+'_'+str(corr[1])+'_'+'.png')    
    if labels:
        plt.ylabel(labels['ylabel'])
        plt.xlabel(labels['xlabel'])
    #plt.show()
    plt.savefig(os.path.join(folderpath,'plt_'+figname+'.pdf'))
    plt.close()
    
def plot_hist(mat1,figname, folderpath,dbins):
    '''
    Need to define what is the bin size
    '''
    plt.hist(mat1, dbins)
    #plt.show()
    plt.savefig(os.path.join(folderpath,'plt_'+figname+'.pdf'))
    plt.close()
    
def plot_hist2(mat1,mat2, figname, folderpath,dbins, labels={}):
    '''
    Need to define what is the bin size
    mat1 all values
    mat2 a subset of mat1 values
    '''
    plt.hist(mat1, dbins,facecolor='blue')
    plt.hold(True)
    plt.hist(mat2, dbins, facecolor='red', alpha=0.5)
    if labels:
        plt.ylabel(labels['ylabel'])
        plt.xlabel(labels['xlabel'])

    #plt.show()
    plt.savefig(os.path.join(folderpath,'plt_'+figname+'.pdf'))
    plt.close()


def plot_hist_ratio(mat1,mat2, figname, folderpath, dbins,labels={}):
    '''
    Need to define what is the bin size
    mat1 all values
    mat2 a subset of mat1 values
    '''
    # Total
    h1, db1 = np.histogram(mat1, dbins)
    # Subset
    h2, db2 = np.histogram(mat2, dbins)
    ratio = np.true_divide(h2,(h1+1))
    
    print dbins[ratio == 0]
    print "The length of the zero ratio is:"
    print len(dbins[ratio == 0])
       
    plt.scatter(dbins[1:], ratio, c='k', s=2, marker='o')
    if labels:
        plt.ylabel(labels['ylabel'])
        plt.xlabel(labels['xlabel'])

    #plt.show()
    plt.savefig(os.path.join(folderpath,'plt_'+figname+'.pdf'))
    plt.close()
    
    plot_hist(ratio,'plt_'+figname+'_hist_', folderpath,len(dbins))
    
    hind2 = np.digitize(mat1, dbins[ratio == 0])
    
    pp=dbins[ratio == 0]
    
    for i,j in enumerate(hind2):
        if i in pp:
            print i,j
    

def plot_2d_histogram(mat1,mat2, dbins,figname, folderpath):
    """
    """
    #A = np.random.randint(10, 100, 100).reshape(10, 10)
    #mask =  np.tri(A.shape[0],A.shape[1], k=-1)
    #A = np.ma.array(A, mask=mask) # mask out the lower triangle
    
    H, xedges, yedges =np.histogram2d(mat1, mat2, bins=dbins, range=None, normed=False, weights=None)
    # Plot the figure
    extent = [yedges[0], yedges[-1], xedges[-1], xedges[0]]
    plt.imshow(H, extent=extent, interpolation='nearest')
    plt.colorbar()
    #plt.show()
    plt.savefig(os.path.join(folderpath,'plt_'+figname+'.pdf'))
    plt.close()

def plot_mat2d(mat2d,xedges, yedges, labels, figname, folderpath):
    '''
    '''
    extent = [yedges[0], yedges[-1], xedges[-1], xedges[0]]
    plt.imshow(mat2d, extent=extent, interpolation='nearest')
    plt.colorbar()
    #plt.show()
    if labels:
        plt.ylabel(labels['ylabel'])
        plt.xlabel(labels['xlabel'])
    
    
    plt.savefig(os.path.join(folderpath,'plt_'+figname+'.pdf'))
    plt.close()

    
